import os
import numpy as np
from PIL import Image
'''将映射后的灰度图片转化为彩色图片，用于检查映射是否正确'''
# 定义标签颜色信息
label_colors = {
    1: (128, 0, 0),  # building - Red
    2: (128, 128, 128),  # sky - Grey
    3: (128, 64, 128),  # road - Pink
    4: (128, 128, 0),  # vegetation - Dark yellow
    5: (0, 0, 192),  # sidewalk/pave - Blue
    6: (64, 0, 128),  # car - Purple
    7: (64, 64, 0),  # pedestrian - Yellow-brown
    8: (0, 128, 192),  # cyclist - Light blue
    9: (192, 128, 128),  # signage - Salmon
    10: (64, 64, 128),  # fence/wall - Grey-purple
    11: (192, 192, 128),  # pole - Light yellow
    0: (0, 0, 0)  # other - Black
}


def grayscale_to_color(input_image_path, output_image_path, label_colors):
    # 打开灰度图像
    img = Image.open(input_image_path)

    # 创建新的彩色图像
    color_img = Image.new('RGB', img.size)

        # 遍历像素并根据标签颜色信息设置彩色像素
        for y in range(img.size[1]):
        for x in range(img.size[0]):
            pixel_value = img.getpixel((x, y))
            color = label_colors[pixel_value]
            color_img.putpixel((x, y), color)

    # 保存彩色图像
    color_img.save(output_image_path)


# 示例用法
data_root = 'D:\guomengqi\毕业倒计时\课题\语义地图处理部分\deeplab_v3\mapping_img_change\data_change0528'
input_image_path = os.path.join(data_root, 'label_2d_img', '000000_bw.png')
output_image_path = os.path.join(data_root, 'test', 'example_color.png')
grayscale_to_color(input_image_path, output_image_path, label_colors)
